import traceback
import json
import logging
import zipfile
import subprocess
import os
import platform
from flask import Flask, request, stream_with_context
from flask_sock import Sock

import utils.html_parser as html_parser
import utils.pre_process as pre_process
from llm.call_ocr import call_ocr_pdf, encode_pdf
from services.code_record_service import CodeRecordGenerationService
from utils.code_parser import Parser
from utils.docx_utils import get_docx_headings_and_formulas
from utils.llm_utils import get_pseudo_code_from_model_res
from llm.call_llm import rule_based_agent
from utils.ocr_utils import get_headings_and_formulas

app = Flask(__name__)
sock = Sock(app)

def convert_docx_to_pdf(docx_path, pdf_path):
    """
    跨平台的docx转pdf函数
    在Windows上使用docx2pdf，在Linux上使用LibreOffice
    """
    try:
        system = platform.system()
        
        if system == "Windows":
            # Windows系统使用docx2pdf
            try:
                from docx2pdf import convert
                convert(docx_path, pdf_path)
                return True
            except ImportError:
                print("docx2pdf not available, falling back to LibreOffice")
                return _convert_with_libreoffice(docx_path, pdf_path)
        else:
            # Linux/Mac系统使用LibreOffice
            return _convert_with_libreoffice(docx_path, pdf_path)
            
    except Exception as e:
        print(f"转换失败: {e}")
        return False

def _convert_with_libreoffice(docx_path, pdf_path):
    """
    使用LibreOffice进行转换
    """
    try:
        # 获取输出目录
        output_dir = os.path.dirname(pdf_path)
        
        # 构造LibreOffice命令，添加字体和编码相关参数
        cmd = [
            'libreoffice',
            '--headless',
            '--convert-to', 'pdf',
            '--outdir', output_dir,
            '--accept=pipe,name=temp;urp;StarOffice.ServiceManager',
            '--nofirststartwizard',
            '--norestore',
            docx_path
        ]
        
        # 设置环境变量以支持中文
        env = os.environ.copy()
        env['LANG'] = 'zh_CN.UTF-8'
        env['LC_ALL'] = 'zh_CN.UTF-8'
        
        # 执行转换
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60, env=env)
        
        if result.returncode == 0:
            # LibreOffice会自动生成与输入文件同名的PDF文件
            generated_pdf = os.path.join(output_dir, os.path.splitext(os.path.basename(docx_path))[0] + '.pdf')
            
            # 如果生成的文件名与期望的不同，则重命名
            if generated_pdf != pdf_path and os.path.exists(generated_pdf):
                os.rename(generated_pdf, pdf_path)
            
            return True
        else:
            print(f"LibreOffice转换失败: {result.stderr}")
            return False
            
    except subprocess.TimeoutExpired:
        print("LibreOffice转换超时")
        return False
    except Exception as e:
        print(f"LibreOffice转换异常: {e}")
        return False

@app.route("/", methods=['GET'])
def hello():
    # to test ping
    return "Hello World!"

@app.route('/get_code', methods=['POST'])
def get_code():
    params = request.json
    codeRecordService = CodeRecordGenerationService()
    module_id = params['moduleId']
    module_name = params['module']
    pseudo_code = params['pseudoCode']
    dataflow_reporting: bool = params['dataflowReporting']
    try:
        config_path='./configs/dicts/config_tmp.json' # 临时保存json文件
        module_dict_str=pre_process.pre_process(params["dict"],True)
        if module_dict_str=="":
            module_dict_json=[{},{},{},{},"数据字典为空"]
        else:
            module_dict_json = html_parser.str_to_json_502_new(module_dict_str)
        project_dict_str=pre_process.pre_process(params["project_dict"])
        if project_dict_str!="":
            project_dict_json=html_parser.str_to_json_502_new(project_dict_str)
            module_dict_json = html_parser.merge_json(project_dict_json,module_dict_json)
            if type(module_dict_json)==str and module_dict_json.startswith("ERROR"):
                return module_dict_json

        # 合并domainKnowledge
        config_dict_str=params["configure_dict"]
        config_dict_str=pre_process.pre_process(config_dict_str,True)
        under_func_res,macro_res,key_word_config=html_parser.domain_to_json_502(config_dict_str)
        under_func_res.update(module_dict_json[1])
        module_dict_json[1]=under_func_res
        for key in macro_res.keys():
            if key in module_dict_json[2].keys():
                return "ERROR:宏常量中有重复的参数名 {}".format(key)
        macro_res.update(module_dict_json[2])
        module_dict_json[2]=macro_res
        key_words=key_word_config
        config_json=module_dict_json

        with open(config_path, 'w',encoding="utf-8") as f:
            json.dump(config_json, f,ensure_ascii=False)
        pseudoCode = pre_process.pre_process(params["pseudoCode"])
        pseudoCode = pre_process.add_semicolon(pseudoCode)
        # print("预处理后的伪代码：\n{}".format(pseudoCode)+"\n"+"-"*50)
        lines = pseudoCode.split("\r\n")
        func_statment = ""
        func_pseudo_code = ""
        res_code = ""
        for line in lines:
            if line.strip().startswith("###") and line.strip().endswith("###;"):
                if func_pseudo_code.strip():
                    parser = Parser(func_pseudo_code,config_path,key_words= key_words)
                    parser.parse()
                    if func_statment != "":
                        res_code += func_statment + "\r\n{"
                    res_code += parser.print_codes() + "\r\n"
                    if func_statment != "":
                        res_code += "return ;\r\n}\r\n\r\n"
                func_pseudo_code = ""
                func_statment = line[line.find("###")+3:line.rfind("###")] 
            else :
                func_pseudo_code += line + "\r\n"
                
        if func_pseudo_code.strip():
            parser = Parser(func_pseudo_code,config_path,key_words= key_words)
            parser.parse()
            if func_statment != "":
                res_code += func_statment + "\r\n{"
            res_code += parser.print_codes() + "\r\n"
            if func_statment != "":
                res_code += "return ;\r\n}\r\n"

        if dataflow_reporting:
            codeRecordService.update_record_c_status(module_id, module_name, pseudo_code, res_code, None)
        return res_code
    except Exception as e:
        c_code_error = str(e) + "\n" + traceback.format_exc()
        print("ERROR: " + c_code_error)
        if dataflow_reporting:
            codeRecordService.update_record_c_status(module_id, module_name, pseudo_code, None, c_code_error)
        return "ERROR:" + str(e)

@app.route('/upload_docxfile', methods=['POST'])
def upload_docxfile():

    try:
        file=request.files['file']
        #将word内容保存后缀再改成zip
        file.save("tmp.zip")
        with zipfile.ZipFile("tmp.zip", 'r') as zip_ref:
            document = list(filter(lambda f: f == 'word/document.xml', zip_ref.namelist()))
            # 打开并读取文件内容
            with zip_ref.open(document[0]) as f:
                content = f.read().decode("utf-8")
                return content
    except Exception as e:
        return "ERROR:" + str(e)

    
@app.route('/get_llm_code_streaming', methods=['POST'])
def get_llm_code_streaming():
    params = request.json
    """
    dict Str: 数据字典
    pseudoCode Str: 伪代码
    format_latex_prompt Str: prompt 用于格式化latex提示
    combine_prompt Str: prompt 用于合并各部分输入生成最终代码
    """

    format_latex_prompt = params['format_latex_prompt'] if 'format_latex_prompt' in params else ""
    combine_prompt = params['prompt'] if 'prompt' in params else ""
    params = request.json

    config_path='./configs/config_tmp_for_llm_code.json' # 临时保存json文件
    module_dict_str=pre_process.pre_process(params["dict"],True)
    if module_dict_str=="":
        module_dict_json=[{},{},{},{},"数据字典为空"]
    else:
        module_dict_json = html_parser.str_to_json_502_new(module_dict_str)
    project_dict_str=pre_process.pre_process(params["project_dict"])
    if project_dict_str!="":
        project_dict_json=html_parser.str_to_json_502_new(project_dict_str)
        module_dict_json = html_parser.merge_json(project_dict_json,module_dict_json) # merge json
        if type(module_dict_json)==str and module_dict_json.startswith("ERROR"):
            return module_dict_json

    # 合并domainKnowledge
    config_dict_str=params["configure_dict"]
    under_func_res,macro_res,key_word_config=html_parser.domain_to_json_502(config_dict_str)
    under_func_res.update(module_dict_json[1])
    module_dict_json[1]=under_func_res
    for key in macro_res.keys():
        if key in module_dict_json[2].keys():
            return "ERROR:宏常量中有重复的参数名 {}".format(key)
    macro_res.update(module_dict_json[2])
    module_dict_json[2]=macro_res

    key_words=key_word_config
    config_json=module_dict_json
    
    # TODO: 可以修改这个内容来构造输入给大模型的数据字典
    table_str = config_dict_str+project_dict_str+module_dict_str

    with open(config_path, 'w',encoding="utf-8") as f:
        json.dump(config_json, f,ensure_ascii=False)
    pseudoCode = pre_process.pre_process(params["pseudoCode"])
    pseudoCode = pre_process.add_semicolon(pseudoCode)
    # print("预处理后的伪代码：\n{}".format(pseudoCode)+"\n"+"-"*50)
    try:
        parser = Parser(pseudoCode,config_path,key_words=key_words, llm_mode=True)
        parser.parse()
        code = parser.print_codes()
    except Exception as e:
        code = "ERROR:" + str(e)

    # 使用第一次的工具输出作为tool_out
    tool_out = code
    try:
        # 构造输入
        agent_inp = {
            "pseudo_code": pseudoCode,
            "table_json": config_json,
            "config_path": config_path,
            "table_str": table_str,
            "format_latex_prompt": format_latex_prompt,
            "combine_prompt": combine_prompt,
            "first_tool_out": tool_out,
            "key_words": key_words
        }
        code, history = rule_based_agent(agent_inp, is_stream=True)
        return app.response_class(stream_with_context(code()))
    except Exception as e:
        return "ERROR:" + str(e)

from llm.call_with_step import qianwen_with_step, model_with_step
progress_messages = [
    "步骤 1/6 ：检查伪代码书写格式规范 ... ", 
    "步骤 2/6 ：检查控制流结构规范 ... ", 
    "步骤 3/6 ：检查数学运算符号与表示方法 ... ", 
    "步骤 4/6 ：检查向量和矩阵的表示方法 ... ", 
    "步骤 5/6 ：检查额外的改写要求 ... ", 
    "步骤 6/6 ：检查额外的改写要求 ..."
    ]
@sock.route('/get_standard_pseudo_ws')
def get_standard_pseudo_ws(ws):
    # print("WebSocket Connected")
    try :
        data = ws.receive()
        req_data = json.loads(data)["data"]

        pseudo_code = req_data["pseudoCode"]
        invocation_method = req_data["invocationMethod"]
        model_name = req_data["modelName"]
        data_dict = json.dumps(req_data["dict"])
        model_url = req_data["modelUrl"]
        dataflow_reporting: bool = req_data['dataflowReporting']
        initial_pseudo_code = pseudo_code

        def send_progress(step):
            ws.send(json.dumps({
                'type': 'progress',
                'data': progress_messages[step]
            }))
            # print(progress_messages[step])

        intermediate_model_responses = []

        for i in range(0, 6):
            send_progress(i)
            if invocation_method == 'public':
                content = qianwen_with_step(pseudo_code, data_dict, initial_pseudo_code, i)
            else:
                content = model_with_step(model_url, model_name, pseudo_code, data_dict, initial_pseudo_code, i)
            current_code = get_pseudo_code_from_model_res(content)
            if current_code.strip():
                pseudo_code = current_code
            intermediate_model_responses.append(content)

        ws.send(json.dumps({
            'type': 'result',
            'data': pseudo_code
        }))

        if dataflow_reporting:
            codeRecordService = CodeRecordGenerationService()
            codeRecordService.create_or_update_record_pseudo_code(req_data, intermediate_model_responses)

    except Exception as e:
        print(e)
        ws.send(json.dumps({
            'type': 'error',
            'data': f"ERROR: {e}"
        }))
    
    finally:
        ws.close()

@app.route('/import_document', methods=['POST'])
def import_document():
    file= request.files['file']
    if not file:
        return "ERROR: No file uploaded", 400
    if not file.filename:
        return "ERROR: No file name provided", 400
    if not file.filename.endswith(('.pdf', '.docx')):
        return "ERROR: Unsupported file type. Only PDF and DOCX are allowed.", 400
    file_type = file.filename.split('.')[-1].lower()
    # 保存文件，如果是word则另存为pdf
    if file_type == "docx":
        file.save(file.filename)
        pdf_path = file.filename[:-5] + ".pdf"
        
        if convert_docx_to_pdf(file.filename, pdf_path):
            os.remove(file.filename)
        else:
            os.remove(file.filename)
            return "ERROR: 文档转换失败", 500
    else:
        pdf_path = file.filename
        file.save(pdf_path)
    
    base64Data = encode_pdf(pdf_path)
    os.remove(pdf_path)

    ocr_response = call_ocr_pdf(base64Data)
    res=get_headings_and_formulas(ocr_response)
    return json.dumps(res)
    if file_type == "pdf":
        ocr_response = call_ocr_pdf(base64Data)
        return json.dumps(get_headings_and_formulas(ocr_response))
    # todo: docx 中的公式需要识别其中的图片，直接使用 base64 调 ocr 识别不出来
    # todo: 以及公式如果是 latex 格式，应该也需要调用 ocr ，目前只能识别出纯文本
    elif file_type == "docx":
        formulas = get_docx_headings_and_formulas(base64Data)
        return json.dumps({"formulas": formulas, "dict": ""})
    return None


if __name__ == '__main__':
    app.logger.setLevel(logging.ERROR)
    logging.getLogger('werkzeug').setLevel(logging.ERROR)
    app.run(host='0.0.0.0', port=9001, debug=False)